{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "293f12fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b1abf89",
   "metadata": {},
   "source": [
    "# Run me in google colab [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/pandas/split-apply-combine/notebook.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/pandas/split-apply-combine/notebook.ipynb)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/pandas/split-apply-combine/notebook.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2b51ad4e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-28T16:32:44.995610Z",
     "start_time": "2024-06-28T16:32:38.286618Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The hamilton.plugins.jupyter_magic extension is already loaded. To reload it, use:\n",
      "  %reload_ext hamilton.plugins.jupyter_magic\n"
     ]
    }
   ],
   "source": [
    "%load_ext hamilton.plugins.jupyter_magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0528c5ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 10.0.1 (20240210.2158)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"1784pt\" height=\"314pt\"\n",
       " viewBox=\"0.00 0.00 1783.85 314.30\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 310.3)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-310.3 1779.85,-310.3 1779.85,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>cluster__legend</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" points=\"24.12,-168.3 24.12,-298.3 108.97,-298.3 108.97,-168.3 24.12,-168.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"66.55\" y=\"-281\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Legend</text>\n",
       "</g>\n",
       "<!-- tax -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>tax</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1585.75,-126.1C1585.75,-126.1 1549.9,-126.1 1549.9,-126.1 1543.9,-126.1 1537.9,-120.1 1537.9,-114.1 1537.9,-114.1 1537.9,-74.5 1537.9,-74.5 1537.9,-68.5 1543.9,-62.5 1549.9,-62.5 1549.9,-62.5 1585.75,-62.5 1585.75,-62.5 1591.75,-62.5 1597.75,-68.5 1597.75,-74.5 1597.75,-74.5 1597.75,-114.1 1597.75,-114.1 1597.75,-120.1 1591.75,-126.1 1585.75,-126.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1558.07\" y=\"-103\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">tax</text>\n",
       "<text text-anchor=\"start\" x=\"1548.7\" y=\"-75\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">Series</text>\n",
       "</g>\n",
       "<!-- final_tax_dataframe -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>final_tax_dataframe</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1763.85,-96.1C1763.85,-96.1 1638.75,-96.1 1638.75,-96.1 1632.75,-96.1 1626.75,-90.1 1626.75,-84.1 1626.75,-84.1 1626.75,-44.5 1626.75,-44.5 1626.75,-38.5 1632.75,-32.5 1638.75,-32.5 1638.75,-32.5 1763.85,-32.5 1763.85,-32.5 1769.85,-32.5 1775.85,-38.5 1775.85,-44.5 1775.85,-44.5 1775.85,-84.1 1775.85,-84.1 1775.85,-90.1 1769.85,-96.1 1763.85,-96.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1637.55\" y=\"-73\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">final_tax_dataframe</text>\n",
       "<text text-anchor=\"start\" x=\"1667.17\" y=\"-45\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- tax&#45;&gt;final_tax_dataframe -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>tax&#45;&gt;final_tax_dataframe</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1598,-87.64C1603.35,-86.42 1609.16,-85.1 1615.22,-83.71\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1615.88,-87.15 1624.85,-81.52 1614.32,-80.33 1615.88,-87.15\"/>\n",
       "</g>\n",
       "<!-- Tax Credit -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>Tax Credit</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1369.05,-290.1C1369.05,-290.1 1305.45,-290.1 1305.45,-290.1 1299.45,-290.1 1293.45,-284.1 1293.45,-278.1 1293.45,-278.1 1293.45,-238.5 1293.45,-238.5 1293.45,-232.5 1299.45,-226.5 1305.45,-226.5 1305.45,-226.5 1369.05,-226.5 1369.05,-226.5 1375.05,-226.5 1381.05,-232.5 1381.05,-238.5 1381.05,-238.5 1381.05,-278.1 1381.05,-278.1 1381.05,-284.1 1375.05,-290.1 1369.05,-290.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1304.25\" y=\"-267\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">Tax Credit</text>\n",
       "<text text-anchor=\"start\" x=\"1318.12\" y=\"-239\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">Series</text>\n",
       "</g>\n",
       "<!-- tax_formula -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>tax_formula</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1496.9,-126.1C1496.9,-126.1 1422.05,-126.1 1422.05,-126.1 1416.05,-126.1 1410.05,-120.1 1410.05,-114.1 1410.05,-114.1 1410.05,-74.5 1410.05,-74.5 1410.05,-68.5 1416.05,-62.5 1422.05,-62.5 1422.05,-62.5 1496.9,-62.5 1496.9,-62.5 1502.9,-62.5 1508.9,-68.5 1508.9,-74.5 1508.9,-74.5 1508.9,-114.1 1508.9,-114.1 1508.9,-120.1 1502.9,-126.1 1496.9,-126.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1420.85\" y=\"-103\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">tax_formula</text>\n",
       "<text text-anchor=\"start\" x=\"1440.35\" y=\"-75\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">Series</text>\n",
       "</g>\n",
       "<!-- Tax Credit&#45;&gt;tax_formula -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>Tax Credit&#45;&gt;tax_formula</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1373.12,-226.33C1375.9,-223.36 1378.58,-220.33 1381.05,-217.3 1401.66,-192.04 1421.19,-160.95 1435.47,-136.34\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1438.38,-138.31 1440.32,-127.89 1432.31,-134.83 1438.38,-138.31\"/>\n",
       "</g>\n",
       "<!-- over_100k_tax -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>over_100k_tax</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M856.12,-126.1C856.12,-126.1 765.52,-126.1 765.52,-126.1 759.52,-126.1 753.52,-120.1 753.52,-114.1 753.52,-114.1 753.52,-74.5 753.52,-74.5 753.52,-68.5 759.52,-62.5 765.52,-62.5 765.52,-62.5 856.12,-62.5 856.12,-62.5 862.12,-62.5 868.12,-68.5 868.12,-74.5 868.12,-74.5 868.12,-114.1 868.12,-114.1 868.12,-120.1 862.12,-126.1 856.12,-126.1\"/>\n",
       "<text text-anchor=\"start\" x=\"764.32\" y=\"-103\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">over_100k_tax</text>\n",
       "<text text-anchor=\"start\" x=\"776.7\" y=\"-75\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- combined_dataframe -->\n",
       "<g id=\"node15\" class=\"node\">\n",
       "<title>combined_dataframe</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1252.45,-167.1C1252.45,-167.1 1117.6,-167.1 1117.6,-167.1 1111.6,-167.1 1105.6,-161.1 1105.6,-155.1 1105.6,-155.1 1105.6,-115.5 1105.6,-115.5 1105.6,-109.5 1111.6,-103.5 1117.6,-103.5 1117.6,-103.5 1252.45,-103.5 1252.45,-103.5 1258.45,-103.5 1264.45,-109.5 1264.45,-115.5 1264.45,-115.5 1264.45,-155.1 1264.45,-155.1 1264.45,-161.1 1258.45,-167.1 1252.45,-167.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1116.4\" y=\"-144\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">combined_dataframe</text>\n",
       "<text text-anchor=\"start\" x=\"1150.9\" y=\"-116\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- over_100k_tax&#45;&gt;combined_dataframe -->\n",
       "<g id=\"edge23\" class=\"edge\">\n",
       "<title>over_100k_tax&#45;&gt;combined_dataframe</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M868.35,-100.53C928.29,-107.13 1023.83,-117.65 1094.12,-125.4\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1093.34,-128.83 1103.67,-126.45 1094.11,-121.87 1093.34,-128.83\"/>\n",
       "</g>\n",
       "<!-- Income -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>Income</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1360.05,-208.1C1360.05,-208.1 1314.45,-208.1 1314.45,-208.1 1308.45,-208.1 1302.45,-202.1 1302.45,-196.1 1302.45,-196.1 1302.45,-156.5 1302.45,-156.5 1302.45,-150.5 1308.45,-144.5 1314.45,-144.5 1314.45,-144.5 1360.05,-144.5 1360.05,-144.5 1366.05,-144.5 1372.05,-150.5 1372.05,-156.5 1372.05,-156.5 1372.05,-196.1 1372.05,-196.1 1372.05,-202.1 1366.05,-208.1 1360.05,-208.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1313.25\" y=\"-185\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">Income</text>\n",
       "<text text-anchor=\"start\" x=\"1318.12\" y=\"-157\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">Series</text>\n",
       "</g>\n",
       "<!-- Income&#45;&gt;tax_formula -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>Income&#45;&gt;tax_formula</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1372.38,-153.02C1381.63,-146.71 1391.89,-139.71 1402.02,-132.81\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1403.71,-135.89 1409.99,-127.37 1399.76,-130.11 1403.71,-135.89\"/>\n",
       "</g>\n",
       "<!-- split_dataframe -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>split_dataframe</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M272.2,-167.1C272.2,-167.1 174.1,-167.1 174.1,-167.1 168.1,-167.1 162.1,-161.1 162.1,-155.1 162.1,-155.1 162.1,-115.5 162.1,-115.5 162.1,-109.5 168.1,-103.5 174.1,-103.5 174.1,-103.5 272.2,-103.5 272.2,-103.5 278.2,-103.5 284.2,-109.5 284.2,-115.5 284.2,-115.5 284.2,-155.1 284.2,-155.1 284.2,-161.1 278.2,-167.1 272.2,-167.1\"/>\n",
       "<text text-anchor=\"start\" x=\"172.9\" y=\"-144\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">split_dataframe</text>\n",
       "<text text-anchor=\"start\" x=\"212.65\" y=\"-116\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">dict</text>\n",
       "</g>\n",
       "<!-- under_100k -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>under_100k</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M405.55,-208.1C405.55,-208.1 332.95,-208.1 332.95,-208.1 326.95,-208.1 320.95,-202.1 320.95,-196.1 320.95,-196.1 320.95,-156.5 320.95,-156.5 320.95,-150.5 326.95,-144.5 332.95,-144.5 332.95,-144.5 405.55,-144.5 405.55,-144.5 411.55,-144.5 417.55,-150.5 417.55,-156.5 417.55,-156.5 417.55,-196.1 417.55,-196.1 417.55,-202.1 411.55,-208.1 405.55,-208.1\"/>\n",
       "<text text-anchor=\"start\" x=\"331.75\" y=\"-185\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">under_100k</text>\n",
       "<text text-anchor=\"start\" x=\"335.12\" y=\"-157\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- split_dataframe&#45;&gt;under_100k -->\n",
       "<g id=\"edge16\" class=\"edge\">\n",
       "<title>split_dataframe&#45;&gt;under_100k</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M284.43,-152.45C292.7,-154.8 301.18,-157.22 309.43,-159.56\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"308.45,-162.92 319.02,-162.29 310.36,-156.19 308.45,-162.92\"/>\n",
       "</g>\n",
       "<!-- over_100k -->\n",
       "<g id=\"node14\" class=\"node\">\n",
       "<title>over_100k</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M402.17,-126.1C402.17,-126.1 336.32,-126.1 336.32,-126.1 330.32,-126.1 324.32,-120.1 324.32,-114.1 324.32,-114.1 324.32,-74.5 324.32,-74.5 324.32,-68.5 330.32,-62.5 336.32,-62.5 336.32,-62.5 402.17,-62.5 402.17,-62.5 408.17,-62.5 414.17,-68.5 414.17,-74.5 414.17,-74.5 414.17,-114.1 414.17,-114.1 414.17,-120.1 408.17,-126.1 402.17,-126.1\"/>\n",
       "<text text-anchor=\"start\" x=\"336.25\" y=\"-103\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">over_100k</text>\n",
       "<text text-anchor=\"start\" x=\"335.12\" y=\"-75\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- split_dataframe&#45;&gt;over_100k -->\n",
       "<g id=\"edge21\" class=\"edge\">\n",
       "<title>split_dataframe&#45;&gt;over_100k</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M284.43,-118.15C293.96,-115.44 303.78,-112.64 313.19,-109.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"313.88,-113.41 322.54,-107.3 311.97,-106.67 313.88,-113.41\"/>\n",
       "</g>\n",
       "<!-- tax_formula&#45;&gt;tax -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>tax_formula&#45;&gt;tax</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1509.29,-94.3C1515,-94.3 1520.77,-94.3 1526.35,-94.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1525.92,-97.8 1535.92,-94.3 1525.92,-90.8 1525.92,-97.8\"/>\n",
       "</g>\n",
       "<!-- tax_formula&#45;&gt;final_tax_dataframe -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>tax_formula&#45;&gt;final_tax_dataframe</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1497.1,-62.02C1516.44,-48.15 1541.39,-34.92 1566.82,-34.3\"/>\n",
       "</g>\n",
       "<!-- under_100k_tax -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>under_100k_tax</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1064.6,-202.1C1064.6,-202.1 965,-202.1 965,-202.1 959,-202.1 953,-196.1 953,-190.1 953,-190.1 953,-150.5 953,-150.5 953,-144.5 959,-138.5 965,-138.5 965,-138.5 1064.6,-138.5 1064.6,-138.5 1070.6,-138.5 1076.6,-144.5 1076.6,-150.5 1076.6,-150.5 1076.6,-190.1 1076.6,-190.1 1076.6,-196.1 1070.6,-202.1 1064.6,-202.1\"/>\n",
       "<text text-anchor=\"start\" x=\"963.8\" y=\"-179\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">under_100k_tax</text>\n",
       "<text text-anchor=\"start\" x=\"980.67\" y=\"-151\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- under_100k_tax&#45;&gt;combined_dataframe -->\n",
       "<g id=\"edge22\" class=\"edge\">\n",
       "<title>under_100k_tax&#45;&gt;combined_dataframe</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1077.01,-157.57C1082.62,-156.4 1088.38,-155.2 1094.17,-153.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1094.77,-157.45 1103.85,-151.98 1093.34,-150.59 1094.77,-157.45\"/>\n",
       "</g>\n",
       "<!-- under_100k_tax.with_tax_credit -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>under_100k_tax.with_tax_credit</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M912,-208.1C912,-208.1 709.65,-208.1 709.65,-208.1 703.65,-208.1 697.65,-202.1 697.65,-196.1 697.65,-196.1 697.65,-156.5 697.65,-156.5 697.65,-150.5 703.65,-144.5 709.65,-144.5 709.65,-144.5 912,-144.5 912,-144.5 918,-144.5 924,-150.5 924,-156.5 924,-156.5 924,-196.1 924,-196.1 924,-202.1 918,-208.1 912,-208.1\"/>\n",
       "<text text-anchor=\"start\" x=\"708.45\" y=\"-185\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">under_100k_tax.with_tax_credit</text>\n",
       "<text text-anchor=\"start\" x=\"776.7\" y=\"-157\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- under_100k_tax.with_tax_credit&#45;&gt;under_100k_tax -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>under_100k_tax.with_tax_credit&#45;&gt;under_100k_tax</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M924.33,-172.96C930.08,-172.79 935.79,-172.62 941.37,-172.45\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"941.21,-175.96 951.1,-172.16 941,-168.96 941.21,-175.96\"/>\n",
       "</g>\n",
       "<!-- Tax Rate -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>Tax Rate</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M1364.17,-126.1C1364.17,-126.1 1310.32,-126.1 1310.32,-126.1 1304.32,-126.1 1298.32,-120.1 1298.32,-114.1 1298.32,-114.1 1298.32,-74.5 1298.32,-74.5 1298.32,-68.5 1304.32,-62.5 1310.32,-62.5 1310.32,-62.5 1364.17,-62.5 1364.17,-62.5 1370.17,-62.5 1376.17,-68.5 1376.17,-74.5 1376.17,-74.5 1376.17,-114.1 1376.17,-114.1 1376.17,-120.1 1370.17,-126.1 1364.17,-126.1\"/>\n",
       "<text text-anchor=\"start\" x=\"1309.12\" y=\"-103\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">Tax Rate</text>\n",
       "<text text-anchor=\"start\" x=\"1318.12\" y=\"-75\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">Series</text>\n",
       "</g>\n",
       "<!-- Tax Rate&#45;&gt;tax_formula -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>Tax Rate&#45;&gt;tax_formula</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1376.62,-94.3C1383.65,-94.3 1391.13,-94.3 1398.61,-94.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1398.39,-97.8 1408.39,-94.3 1398.39,-90.8 1398.39,-97.8\"/>\n",
       "</g>\n",
       "<!-- under_100k_tax.with_tax_rate -->\n",
       "<g id=\"node13\" class=\"node\">\n",
       "<title>under_100k_tax.with_tax_rate</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M656.65,-208.1C656.65,-208.1 466.3,-208.1 466.3,-208.1 460.3,-208.1 454.3,-202.1 454.3,-196.1 454.3,-196.1 454.3,-156.5 454.3,-156.5 454.3,-150.5 460.3,-144.5 466.3,-144.5 466.3,-144.5 656.65,-144.5 656.65,-144.5 662.65,-144.5 668.65,-150.5 668.65,-156.5 668.65,-156.5 668.65,-196.1 668.65,-196.1 668.65,-202.1 662.65,-208.1 656.65,-208.1\"/>\n",
       "<text text-anchor=\"start\" x=\"465.1\" y=\"-185\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">under_100k_tax.with_tax_rate</text>\n",
       "<text text-anchor=\"start\" x=\"527.35\" y=\"-157\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- under_100k&#45;&gt;under_100k_tax.with_tax_rate -->\n",
       "<g id=\"edge19\" class=\"edge\">\n",
       "<title>under_100k&#45;&gt;under_100k_tax.with_tax_rate</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M417.76,-176.3C425.57,-176.3 433.95,-176.3 442.59,-176.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"442.58,-179.8 452.58,-176.3 442.58,-172.8 442.58,-179.8\"/>\n",
       "</g>\n",
       "<!-- over_100k_tax.with_tax_rate -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>over_100k_tax.with_tax_rate</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M652.15,-126.1C652.15,-126.1 470.8,-126.1 470.8,-126.1 464.8,-126.1 458.8,-120.1 458.8,-114.1 458.8,-114.1 458.8,-74.5 458.8,-74.5 458.8,-68.5 464.8,-62.5 470.8,-62.5 470.8,-62.5 652.15,-62.5 652.15,-62.5 658.15,-62.5 664.15,-68.5 664.15,-74.5 664.15,-74.5 664.15,-114.1 664.15,-114.1 664.15,-120.1 658.15,-126.1 652.15,-126.1\"/>\n",
       "<text text-anchor=\"start\" x=\"469.6\" y=\"-103\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">over_100k_tax.with_tax_rate</text>\n",
       "<text text-anchor=\"start\" x=\"527.35\" y=\"-75\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- over_100k_tax.with_tax_rate&#45;&gt;over_100k_tax -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>over_100k_tax.with_tax_rate&#45;&gt;over_100k_tax</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M664.51,-94.3C690.58,-94.3 717.98,-94.3 741.88,-94.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"741.79,-97.8 751.79,-94.3 741.79,-90.8 741.79,-97.8\"/>\n",
       "</g>\n",
       "<!-- under_100k_tax.with_tax_rate&#45;&gt;under_100k_tax.with_tax_credit -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>under_100k_tax.with_tax_rate&#45;&gt;under_100k_tax.with_tax_credit</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M668.81,-176.3C674.55,-176.3 680.35,-176.3 686.15,-176.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"685.76,-179.8 695.76,-176.3 685.76,-172.8 685.76,-179.8\"/>\n",
       "</g>\n",
       "<!-- over_100k&#45;&gt;over_100k_tax.with_tax_rate -->\n",
       "<g id=\"edge17\" class=\"edge\">\n",
       "<title>over_100k&#45;&gt;over_100k_tax.with_tax_rate</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M414.37,-94.3C424.52,-94.3 435.76,-94.3 447.38,-94.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"447.1,-97.8 457.1,-94.3 447.1,-90.8 447.1,-97.8\"/>\n",
       "</g>\n",
       "<!-- combined_dataframe&#45;&gt;Tax Credit -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>combined_dataframe&#45;&gt;Tax Credit</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1225.1,-167.3C1244.34,-183.04 1267.7,-202.18 1288.04,-218.83\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1285.72,-221.45 1295.68,-225.08 1290.15,-216.04 1285.72,-221.45\"/>\n",
       "</g>\n",
       "<!-- combined_dataframe&#45;&gt;Income -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>combined_dataframe&#45;&gt;Income</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1264.66,-156.76C1273.74,-159.24 1282.74,-161.7 1291.15,-163.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1290.17,-167.35 1300.74,-166.61 1292.01,-160.6 1290.17,-167.35\"/>\n",
       "</g>\n",
       "<!-- combined_dataframe&#45;&gt;final_tax_dataframe -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>combined_dataframe&#45;&gt;final_tax_dataframe</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1568.82,-34.3C1584.08,-33.93 1600.09,-35.69 1615.48,-38.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1614.58,-41.91 1625.08,-40.46 1615.97,-35.05 1614.58,-41.91\"/>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1217.74,-103.08C1237.65,-85.04 1264.91,-64.06 1293.45,-53.3 1407.41,-10.32 1445.07,-37.29 1566.82,-34.3\"/>\n",
       "</g>\n",
       "<!-- combined_dataframe&#45;&gt;Tax Rate -->\n",
       "<g id=\"edge15\" class=\"edge\">\n",
       "<title>combined_dataframe&#45;&gt;Tax Rate</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1264.66,-113.84C1272.3,-111.75 1279.89,-109.68 1287.12,-107.71\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1287.94,-111.11 1296.66,-105.1 1286.09,-104.36 1287.94,-111.11\"/>\n",
       "</g>\n",
       "<!-- _split_dataframe_inputs -->\n",
       "<g id=\"node16\" class=\"node\">\n",
       "<title>_split_dataframe_inputs</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"133.1,-157.6 0,-157.6 0,-113 133.1,-113 133.1,-157.6\"/>\n",
       "<text text-anchor=\"start\" x=\"14.8\" y=\"-129.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input</text>\n",
       "<text text-anchor=\"start\" x=\"50.05\" y=\"-129.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- _split_dataframe_inputs&#45;&gt;split_dataframe -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>_split_dataframe_inputs&#45;&gt;split_dataframe</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M133.56,-135.3C139.12,-135.3 144.75,-135.3 150.35,-135.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"150.23,-138.8 160.23,-135.3 150.23,-131.8 150.23,-138.8\"/>\n",
       "</g>\n",
       "<!-- _under_100k_tax.with_tax_credit_inputs -->\n",
       "<g id=\"node17\" class=\"node\">\n",
       "<title>_under_100k_tax.with_tax_credit_inputs</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"622.4,-270.6 500.55,-270.6 500.55,-226 622.4,-226 622.4,-270.6\"/>\n",
       "<text text-anchor=\"start\" x=\"515.35\" y=\"-242.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">tax_credits</text>\n",
       "<text text-anchor=\"start\" x=\"586.6\" y=\"-242.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">dict</text>\n",
       "</g>\n",
       "<!-- _under_100k_tax.with_tax_credit_inputs&#45;&gt;under_100k_tax.with_tax_credit -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>_under_100k_tax.with_tax_credit_inputs&#45;&gt;under_100k_tax.with_tax_credit</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M622.67,-230.78C642.54,-224.99 665.42,-218.33 688.2,-211.7\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"688.93,-215.14 697.56,-208.98 686.98,-208.41 688.93,-215.14\"/>\n",
       "</g>\n",
       "<!-- _over_100k_tax.with_tax_rate_inputs -->\n",
       "<g id=\"node18\" class=\"node\">\n",
       "<title>_over_100k_tax.with_tax_rate_inputs</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"425.3,-44.6 313.2,-44.6 313.2,0 425.3,0 425.3,-44.6\"/>\n",
       "<text text-anchor=\"start\" x=\"328\" y=\"-16.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">tax_rates</text>\n",
       "<text text-anchor=\"start\" x=\"389.5\" y=\"-16.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">dict</text>\n",
       "</g>\n",
       "<!-- _over_100k_tax.with_tax_rate_inputs&#45;&gt;over_100k_tax.with_tax_rate -->\n",
       "<g id=\"edge18\" class=\"edge\">\n",
       "<title>_over_100k_tax.with_tax_rate_inputs&#45;&gt;over_100k_tax.with_tax_rate</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M425.78,-43.32C437.98,-47.94 451.24,-52.96 464.52,-57.98\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"463.15,-61.21 473.74,-61.47 465.63,-54.66 463.15,-61.21\"/>\n",
       "</g>\n",
       "<!-- _under_100k_tax.with_tax_rate_inputs -->\n",
       "<g id=\"node19\" class=\"node\">\n",
       "<title>_under_100k_tax.with_tax_rate_inputs</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"425.3,-270.6 313.2,-270.6 313.2,-226 425.3,-226 425.3,-270.6\"/>\n",
       "<text text-anchor=\"start\" x=\"328\" y=\"-242.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">tax_rates</text>\n",
       "<text text-anchor=\"start\" x=\"389.5\" y=\"-242.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">dict</text>\n",
       "</g>\n",
       "<!-- _under_100k_tax.with_tax_rate_inputs&#45;&gt;under_100k_tax.with_tax_rate -->\n",
       "<g id=\"edge20\" class=\"edge\">\n",
       "<title>_under_100k_tax.with_tax_rate_inputs&#45;&gt;under_100k_tax.with_tax_rate</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M425.78,-227.28C437.98,-222.66 451.24,-217.64 464.52,-212.62\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"465.63,-215.94 473.74,-209.13 463.15,-209.39 465.63,-215.94\"/>\n",
       "</g>\n",
       "<!-- input -->\n",
       "<g id=\"node20\" class=\"node\">\n",
       "<title>input</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"93.55,-212.6 39.55,-212.6 39.55,-176 93.55,-176 93.55,-212.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"66.55\" y=\"-188.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input</text>\n",
       "</g>\n",
       "<!-- function -->\n",
       "<g id=\"node21\" class=\"node\">\n",
       "<title>function</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M88.97,-267.6C88.97,-267.6 44.12,-267.6 44.12,-267.6 38.12,-267.6 32.12,-261.6 32.12,-255.6 32.12,-255.6 32.12,-243 32.12,-243 32.12,-237 38.12,-231 44.12,-231 44.12,-231 88.97,-231 88.97,-231 94.97,-231 100.97,-237 100.97,-243 100.97,-243 100.97,-255.6 100.97,-255.6 100.97,-261.6 94.97,-267.6 88.97,-267.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"66.55\" y=\"-243.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">function</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x1484b22b0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%cell_to_module my_functions --display\n",
    "\n",
    "from typing import Dict\n",
    "\n",
    "import numpy as np\n",
    "import pandas\n",
    "import pandas as pd\n",
    "from pandas import DataFrame, Series\n",
    "\n",
    "from hamilton.function_modifiers import extract_columns, extract_fields, inject, pipe, source, step\n",
    "\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "# Tax calculation private functions\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def _tax_rate(df: DataFrame, tax_rates: Dict[str, float]) -> DataFrame:\n",
    "    \"\"\"\n",
    "    Add a series 'Tax Rate' to the DataFrame based on the tax_rates rules.\n",
    "    :param df: The DataFrame\n",
    "    :param tax_rates: Tax rates rules\n",
    "    :return: the DataFrame with the 'Tax Rate' Series\n",
    "    \"\"\"\n",
    "    output = DataFrame()\n",
    "    for tax_rate_formula, tax_rate in tax_rates.items():\n",
    "        selected = df.query(tax_rate_formula)\n",
    "        if selected.empty:\n",
    "            continue\n",
    "        tmp = DataFrame({\"Tax Rate\": tax_rate}, index=selected.index)\n",
    "        output = pd.concat([output, tmp], axis=0)\n",
    "    df = pd.concat([df, output], axis=1)\n",
    "    return df\n",
    "\n",
    "\n",
    "def _tax_credit(df: DataFrame, tax_credits: Dict[str, float]) -> DataFrame:\n",
    "    \"\"\"\n",
    "    Add a series 'Tax Credit' to the DataFrame based on the tax_credits rules.\n",
    "    :param df: The DataFrame\n",
    "    :param tax_credits: Tax credits rules\n",
    "    :return: the DataFrame with the 'Tax Credit' Series\n",
    "    \"\"\"\n",
    "    output = DataFrame()\n",
    "    for tax_credit_formula, tax_credit in tax_credits.items():\n",
    "        selected = df.query(tax_credit_formula)\n",
    "        if selected.empty:\n",
    "            continue\n",
    "        tmp = DataFrame({\"Tax Credit\": tax_credit}, index=selected.index)\n",
    "        output = pd.concat([output, tmp], axis=0)\n",
    "    df = pd.concat([df, output], axis=1)\n",
    "    return df\n",
    "\n",
    "\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "# DataFlow: The functions defined below are displayed in the order of execution\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "@extract_fields({\"under_100k\": DataFrame, \"over_100k\": DataFrame})\n",
    "# Step 1: DataFrame is split in 2 DataFrames\n",
    "def split_dataframe(input: DataFrame) -> Dict[str, DataFrame]:\n",
    "    \"\"\"\n",
    "    That function takes the DataFrame in input and split it in 2 DataFrames:\n",
    "      - under_100k: Rows where 'Income' is under 100k\n",
    "      - over_100k: Rows where 'Income' is over 100k\n",
    "\n",
    "    :param input: the DataFrame to process\n",
    "    :return: a Dict with the DataFrames and the Tax Rates & Credit rules\n",
    "    \"\"\"\n",
    "    return {\n",
    "        \"under_100k\": input.query(\"Income < 100000\"),\n",
    "        \"over_100k\": input.query(\"Income >= 100000\"),\n",
    "    }\n",
    "\n",
    "\n",
    "@pipe(\n",
    "    step(_tax_rate, tax_rates=source(\"tax_rates\")),  # apply the _tax_rate step\n",
    "    step(_tax_credit, tax_credits=source(\"tax_credits\")),  # apply the _tax_credit step\n",
    ")\n",
    "# Step 2: DataFrame for Income under 100k applies a tax calculation pipeline\n",
    "def under_100k_tax(under_100k: DataFrame) -> DataFrame:\n",
    "    \"\"\"\n",
    "    Tax calculation pipeline for 'Income' under 100k.\n",
    "    :param under_100k: The DataFrame  where 'Income' is under 100k\n",
    "    :return: the DataFrame with the 'Tax' Series\n",
    "    \"\"\"\n",
    "    return under_100k\n",
    "\n",
    "\n",
    "@pipe(\n",
    "    step(_tax_rate, tax_rates=source(\"tax_rates\")),  # apply the _tax_rate step\n",
    ")\n",
    "# Step 2: DataFrame for Income over 100k applies a tax calculation pipeline\n",
    "def over_100k_tax(over_100k: DataFrame) -> DataFrame:\n",
    "    \"\"\"\n",
    "    Tax calculation pipeline for 'Income' over 100k.\n",
    "    :param over_100k: The DataFrame where 'Income' is over 100k\n",
    "    :return: the DataFrame with the 'Tax' Series\n",
    "    \"\"\"\n",
    "    return over_100k\n",
    "\n",
    "\n",
    "@extract_columns(\"Income\", \"Tax Rate\", \"Tax Credit\")\n",
    "# Step 3: DataFrames are combined. Series 'Income', 'Tax Rate', 'Tax Credit' are extracted for next processing step\n",
    "def combined_dataframe(under_100k_tax: DataFrame, over_100k_tax: DataFrame) -> DataFrame:\n",
    "    \"\"\"\n",
    "    That function combine the DataFrames under_100k_tax and over_100k_tax\n",
    "\n",
    "    The @extract_columns decorator is making the Series available for processing.\n",
    "    \"\"\"\n",
    "    combined = pd.concat([under_100k_tax, over_100k_tax], axis=0).sort_index()\n",
    "    return combined\n",
    "\n",
    "\n",
    "# We use @inject decorator here because we have spaces in the names of columns.\n",
    "# If column names are valid python variable names we wouldn't need this.\n",
    "@inject(income=source(\"Income\"), tax_rate=source(\"Tax Rate\"), tax_credit=source(\"Tax Credit\"))\n",
    "# Step 4: 'Tax Formula' is calculated from 'Income', 'Tax Rate' and 'Tax Credit' series\n",
    "def tax_formula(income: Series, tax_rate: Series, tax_credit: Series) -> Series:\n",
    "    \"\"\"\n",
    "    Return a DataFrame with a series 'Tax Formula' from 'Income', 'Tax Rate' and 'Tax Credit' series.\n",
    "\n",
    "    :param income: the 'Income' series\n",
    "    :param tax_rate: the 'Tax Rate' series\n",
    "    :param tax_credit: the 'Tax Credit' series\n",
    "\n",
    "    :return: the DataFrame with the 'Tax Formula' Series\n",
    "    \"\"\"\n",
    "    df = DataFrame({\"income\": income, \"tax_rate\": tax_rate, \"tax_credit\": tax_credit})\n",
    "    df[\"Tax Formula\"] = df.apply(\n",
    "        lambda x: (\n",
    "            f\"({int(x['income'])} * {x['tax_rate']})\"\n",
    "            if np.isnan(x[\"tax_credit\"])\n",
    "            else f\"({int(x['income'])} * {x['tax_rate']}) - ({int(x['income'])} * {x['tax_rate']}) * {x['tax_credit']}\"\n",
    "        ),\n",
    "        axis=1,\n",
    "    )\n",
    "    return df[\"Tax Formula\"]\n",
    "\n",
    "\n",
    "# Step 5: 'Tax' is calculated from 'Tax Formula' series\n",
    "def tax(tax_formula: Series) -> Series:\n",
    "    \"\"\"\n",
    "    Return a series 'Tax' from 'Tax Formula' series.\n",
    "    :param tax_formula: the 'Tax Formula' series.\n",
    "    :return: the 'Tax Formula' Series\n",
    "    \"\"\"\n",
    "    df = tax_formula.to_frame()\n",
    "    df[\"Tax\"] = df[\"Tax Formula\"].apply(lambda x: round(pandas.eval(x)))\n",
    "    return df[\"Tax\"]\n",
    "\n",
    "\n",
    "# Step 6 (final): DataFrame and Series computed are combined\n",
    "def final_tax_dataframe(\n",
    "    combined_dataframe: DataFrame, tax_formula: Series, tax: Series\n",
    ") -> DataFrame:\n",
    "    \"\"\"\n",
    "    That function combine the DataFrame and the 'Tax' and 'Tax Formula' series\n",
    "    \"\"\"\n",
    "    df = combined_dataframe.copy(deep=True)\n",
    "\n",
    "    # Set the 'Tax' and 'Tax Formula' series\n",
    "    df[\"Tax Formula\"] = tax_formula\n",
    "    df[\"Tax\"] = tax\n",
    "\n",
    "    # Transform  the 'Tax Rate' and 'Tax Credit' series to display percentage\n",
    "    df[\"Tax Rate\"] = df[\"Tax Rate\"].apply(lambda x: f\"{int(x * 100)} %\")\n",
    "    df[\"Tax Credit\"] = df[\"Tax Credit\"].apply(\n",
    "        lambda x: f\"{int(x * 100)} %\" if not np.isnan(x) else \"\"\n",
    "    )\n",
    "\n",
    "    # Define the order the DataFrame will be displayed\n",
    "    order = [\"Name\", \"Income\", \"Children\", \"Tax Rate\", \"Tax Credit\", \"Tax\", \"Tax Formula\"]\n",
    "\n",
    "    return df.reindex(columns=order)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bd4d1b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "\n",
    "# import my_functions # this is imported by the cell above\n",
    "from pandas import DataFrame\n",
    "\n",
    "from hamilton import base, driver, lifecycle\n",
    "\n",
    "## if using the Hamilton UI do `pip install sf-hamilton[ui,sdk]`\n",
    "# from hamilton_sdk import adapters\n",
    "\n",
    "# tracker = adapters.HamiltonTracker(\n",
    "#    project_id=4,  # modify this as needed\n",
    "#    username=\"elijah@dagworks.io\",\n",
    "#    dag_name=\"split-apply-combine\",\n",
    "#    tags={\"environment\": \"DEV\", \"team\": \"MY_TEAM\", \"version\": \"1\"}\n",
    "# )\n",
    "\n",
    "driver = (\n",
    "    driver.Builder()\n",
    "    .with_config({})\n",
    "    .with_modules(my_functions)\n",
    "    .with_adapters(\n",
    "        # tracker,  # add tracker if you have the UI set up.\n",
    "        # this is a strict type checker for the input and output of each function.\n",
    "        lifecycle.FunctionInputOutputTypeChecker(),\n",
    "        # this will make execute return a pandas dataframe as a result\n",
    "        base.PandasDataFrameResult(),\n",
    "       \n",
    "    )\n",
    "    .build()\n",
    ")\n",
    "\n",
    "\n",
    "class TaxCalculator:\n",
    "    \"\"\"\n",
    "    Simple class to wrap Hamilton Driver\n",
    "    \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def calculate(\n",
    "        input: DataFrame, tax_rates: Dict[str, float], tax_credits: Dict[str, float]\n",
    "    ) -> DataFrame:\n",
    "        return driver.execute(\n",
    "            inputs={\"input\": input, \"tax_rates\": tax_rates, \"tax_credits\": tax_credits},\n",
    "            final_vars=[\"final_tax_dataframe\"],\n",
    "        )\n",
    "\n",
    "    @staticmethod\n",
    "    def visualize():\n",
    "        # To visualize do `pip install \"sf-hamilton[visualization]\"` if you want these to work\n",
    "        return driver.display_all_functions()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5f8d2fee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         Name  Income  Children Tax Rate Tax Credit    Tax                             Tax Formula\n",
      "0   John        75600         2     20 %        4 %  14515    (75600 * 0.2) - (75600 * 0.2) * 0.04\n",
      "1   Bob         34000         1     15 %        2 %   4998  (34000 * 0.15) - (34000 * 0.15) * 0.02\n",
      "2   Chloe      111500         3     22 %             24530                         (111500 * 0.22)\n",
      "3   Thomas     234546         1     28 %             65673                         (234546 * 0.28)\n",
      "4   Ellis      144865         2     25 %             36216                         (144865 * 0.25)\n",
      "5   Deane      138500         4     25 %             34625                         (138500 * 0.25)\n",
      "6   Mariella    69412         5     18 %       10 %  11245   (69412 * 0.18) - (69412 * 0.18) * 0.1\n",
      "7   Carlos      65535         0     18 %        0 %  11796   (65535 * 0.18) - (65535 * 0.18) * 0.0\n",
      "8   Toney       43642         3     15 %        6 %   6154  (43642 * 0.15) - (43642 * 0.15) * 0.06\n",
      "9   Ramiro     117850         2     22 %             25927                         (117850 * 0.22)\n"
     ]
    }
   ],
   "source": [
    "from inspect import cleandoc\n",
    "from io import StringIO\n",
    "\n",
    "import pandas as pd\n",
    "from pandas import DataFrame\n",
    "\n",
    "\n",
    "def read_table(table: str, delimiter=\"|\") -> DataFrame:\n",
    "    \"\"\"\n",
    "    Read table from string and return pandas DataFrame.\n",
    "    \"\"\"\n",
    "    df = pd.read_table(StringIO(cleandoc(table)), delimiter=delimiter)\n",
    "    df = df.loc[:, ~df.columns.str.match(\"Unnamed\")]\n",
    "    df.columns = df.columns.str.strip()\n",
    "    return df\n",
    "\n",
    "\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "# The Data to process\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "input = read_table(\n",
    "    \"\"\"\n",
    "    | Name     | Income | Children |\n",
    "    | John     | 75600  | 2        |\n",
    "    | Bob      | 34000  | 1        |\n",
    "    | Chloe    | 111500 | 3        |\n",
    "    | Thomas   | 234546 | 1        |\n",
    "    | Ellis    | 144865 | 2        |\n",
    "    | Deane    | 138500 | 4        |\n",
    "    | Mariella | 69412  | 5        |\n",
    "    | Carlos   | 65535  | 0        |\n",
    "    | Toney    | 43642  | 3        |\n",
    "    | Ramiro   | 117850 | 2        |\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "# Tax Rate & Credit rules\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "tax_rates = {\n",
    "    \"Income < 50000\": 0.15,  # < 50k: Tax rate is 15 %\n",
    "    \"Income > 50000 and Income < 70000\": 0.18,  # 50k to 70k: Tax rate is 18 %\n",
    "    \"Income > 70000 and Income < 100000\": 0.2,  # 70k to 100k: Tax rate is 20 %\n",
    "    \"Income > 100000 and Income < 120000\": 0.22,  # 100k to 120k: Tax rate is 22 %\n",
    "    \"Income > 120000 and Income < 150000\": 0.25,  # 120k to 150k: Tax rate is 25 %\n",
    "    \"Income > 150000\": 0.28,  # over 150k: Tax rate is 28 %\n",
    "}\n",
    "\n",
    "tax_credits = {\n",
    "    \"Children == 0\": 0.0,  # 0 child: Tax credit 0 %\n",
    "    \"Children == 1\": 0.02,  # 1 child: Tax credit 2 %\n",
    "    \"Children == 2\": 0.04,  # 2 children: Tax credit 4 %\n",
    "    \"Children == 3\": 0.06,  # 3 children: Tax credit 6 %\n",
    "    \"Children == 4\": 0.08,  # 4 children: Tax credit 8 %\n",
    "    \"Children > 4\": 0.1,  # over 4 children: Tax credit 10 %\n",
    "}\n",
    "\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "# Run the Tax Calculator\n",
    "# ----------------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "# Calculate the taxes\n",
    "output = TaxCalculator.calculate(input, tax_rates, tax_credits)\n",
    "print(output.to_string())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f8215111",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Name</th>\n",
       "      <th>Income</th>\n",
       "      <th>Children</th>\n",
       "      <th>Tax Rate</th>\n",
       "      <th>Tax Credit</th>\n",
       "      <th>Tax</th>\n",
       "      <th>Tax Formula</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>John</td>\n",
       "      <td>75600</td>\n",
       "      <td>2</td>\n",
       "      <td>20 %</td>\n",
       "      <td>4 %</td>\n",
       "      <td>14515</td>\n",
       "      <td>(75600 * 0.2) - (75600 * 0.2) * 0.04</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Bob</td>\n",
       "      <td>34000</td>\n",
       "      <td>1</td>\n",
       "      <td>15 %</td>\n",
       "      <td>2 %</td>\n",
       "      <td>4998</td>\n",
       "      <td>(34000 * 0.15) - (34000 * 0.15) * 0.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Chloe</td>\n",
       "      <td>111500</td>\n",
       "      <td>3</td>\n",
       "      <td>22 %</td>\n",
       "      <td></td>\n",
       "      <td>24530</td>\n",
       "      <td>(111500 * 0.22)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Thomas</td>\n",
       "      <td>234546</td>\n",
       "      <td>1</td>\n",
       "      <td>28 %</td>\n",
       "      <td></td>\n",
       "      <td>65673</td>\n",
       "      <td>(234546 * 0.28)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Ellis</td>\n",
       "      <td>144865</td>\n",
       "      <td>2</td>\n",
       "      <td>25 %</td>\n",
       "      <td></td>\n",
       "      <td>36216</td>\n",
       "      <td>(144865 * 0.25)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Deane</td>\n",
       "      <td>138500</td>\n",
       "      <td>4</td>\n",
       "      <td>25 %</td>\n",
       "      <td></td>\n",
       "      <td>34625</td>\n",
       "      <td>(138500 * 0.25)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Mariella</td>\n",
       "      <td>69412</td>\n",
       "      <td>5</td>\n",
       "      <td>18 %</td>\n",
       "      <td>10 %</td>\n",
       "      <td>11245</td>\n",
       "      <td>(69412 * 0.18) - (69412 * 0.18) * 0.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Carlos</td>\n",
       "      <td>65535</td>\n",
       "      <td>0</td>\n",
       "      <td>18 %</td>\n",
       "      <td>0 %</td>\n",
       "      <td>11796</td>\n",
       "      <td>(65535 * 0.18) - (65535 * 0.18) * 0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Toney</td>\n",
       "      <td>43642</td>\n",
       "      <td>3</td>\n",
       "      <td>15 %</td>\n",
       "      <td>6 %</td>\n",
       "      <td>6154</td>\n",
       "      <td>(43642 * 0.15) - (43642 * 0.15) * 0.06</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Ramiro</td>\n",
       "      <td>117850</td>\n",
       "      <td>2</td>\n",
       "      <td>22 %</td>\n",
       "      <td></td>\n",
       "      <td>25927</td>\n",
       "      <td>(117850 * 0.22)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         Name  Income  Children Tax Rate Tax Credit    Tax  \\\n",
       "0   John        75600         2     20 %        4 %  14515   \n",
       "1   Bob         34000         1     15 %        2 %   4998   \n",
       "2   Chloe      111500         3     22 %             24530   \n",
       "3   Thomas     234546         1     28 %             65673   \n",
       "4   Ellis      144865         2     25 %             36216   \n",
       "5   Deane      138500         4     25 %             34625   \n",
       "6   Mariella    69412         5     18 %       10 %  11245   \n",
       "7   Carlos      65535         0     18 %        0 %  11796   \n",
       "8   Toney       43642         3     15 %        6 %   6154   \n",
       "9   Ramiro     117850         2     22 %             25927   \n",
       "\n",
       "                              Tax Formula  \n",
       "0    (75600 * 0.2) - (75600 * 0.2) * 0.04  \n",
       "1  (34000 * 0.15) - (34000 * 0.15) * 0.02  \n",
       "2                         (111500 * 0.22)  \n",
       "3                         (234546 * 0.28)  \n",
       "4                         (144865 * 0.25)  \n",
       "5                         (138500 * 0.25)  \n",
       "6   (69412 * 0.18) - (69412 * 0.18) * 0.1  \n",
       "7   (65535 * 0.18) - (65535 * 0.18) * 0.0  \n",
       "8  (43642 * 0.15) - (43642 * 0.15) * 0.06  \n",
       "9                         (117850 * 0.22)  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "driver.execute(\n",
    "            inputs={\"input\": input, \"tax_rates\": tax_rates, \"tax_credits\": tax_credits},\n",
    "            final_vars=[\"final_tax_dataframe\"],\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "669eb270",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
